# src/utils.py
import logging
import os
import json
from tqdm import tqdm
import re
import shutil  # 用于删除目录树 (测试时清理) 和复制文件对象
import rarfile  # 用于解压RAR文件
import zipfile  # 用于解压ZIP文件

# 从config导入日志目录和默认级别
try:
    from config import LOGS_DIR, DEFAULT_LOG_LEVEL, PROJECT_ROOT
except ImportError:
    # 如果直接运行utils.py或在测试环境中，config可能无法直接导入
    # 提供一个默认值，或者根据需要调整
    PROJECT_ROOT_FALLBACK = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    LOGS_DIR = os.path.join(PROJECT_ROOT_FALLBACK, "outputs", "logs")
    DEFAULT_LOG_LEVEL = "INFO"
    os.makedirs(LOGS_DIR, exist_ok=True)
    PROJECT_ROOT = PROJECT_ROOT_FALLBACK  # 定义 PROJECT_ROOT 以便测试代码使用


def setup_logger(name, log_file_name, level_str=DEFAULT_LOG_LEVEL):
    """
    配置和返回一个logger实例。

    Args:
        name (str): logger的名称。
        log_file_name (str): 日志文件的名称 (例如 "step1.log")。
        level_str (str): 日志级别字符串 (例如 "INFO", "DEBUG")。

    Returns:
        logging.Logger: 配置好的logger实例。
    """
    log_file_path = os.path.join(LOGS_DIR, log_file_name)
    level = getattr(logging, level_str.upper(), logging.INFO)

    logger = logging.getLogger(name)
    if logger.hasHandlers():  # 防止重复添加handler
        logger.handlers.clear()

    logger.setLevel(level)

    # 创建文件处理器
    fh = logging.FileHandler(log_file_path, encoding='utf-8')
    fh.setLevel(level)

    # 创建控制台处理器
    ch = logging.StreamHandler()
    ch.setLevel(level)

    # 定义输出格式
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    fh.setFormatter(formatter)
    ch.setFormatter(formatter)

    logger.addHandler(fh)
    logger.addHandler(ch)

    return logger


# 通用日志记录器
general_logger = setup_logger("general_utils", "general_utils.log")


def load_text_file(filepath, encoding='utf-8'):
    """逐行读取文本文件内容，忽略空行。"""
    try:
        with open(filepath, 'r', encoding=encoding) as f:
            return [line.strip() for line in f if line.strip()]
    except FileNotFoundError:
        general_logger.error(f"文件未找到: {filepath}")
        return []
    except Exception as e:
        general_logger.error(f"读取文件失败 {filepath}: {e}")
        return []


def save_text_list(data_list, filepath, encoding='utf-8'):
    """将字符串列表保存到文本文件，每行一个字符串。"""
    try:
        with open(filepath, 'w', encoding=encoding) as f:
            for item in data_list:
                f.write(str(item) + '\n')
        general_logger.info(f"数据已保存到: {filepath}")
    except Exception as e:
        general_logger.error(f"保存文件失败 {filepath}: {e}")


def load_json_file(filepath, encoding='utf-8'):
    """加载JSON文件。"""
    try:
        with open(filepath, 'r', encoding=encoding) as f:
            return json.load(f)
    except FileNotFoundError:
        general_logger.error(f"JSON文件未找到: {filepath}")
        return None
    except json.JSONDecodeError as jde:
        general_logger.error(f"JSON文件解码失败: {filepath}. 错误: {jde}")
        return None
    except Exception as e:
        general_logger.error(f"读取JSON文件失败 {filepath}: {e}")
        return None


def save_json_file(data, filepath, encoding='utf-8', indent=4):
    """将数据保存为JSON文件。"""
    try:
        with open(filepath, 'w', encoding=encoding) as f:
            json.dump(data, f, ensure_ascii=False, indent=indent)
        general_logger.info(f"JSON数据已保存到: {filepath}")
    except Exception as e:
        general_logger.error(f"保存JSON文件失败 {filepath}: {e}")


def get_all_files(directory, extensions=('.txt',)):
    """
    递归获取指定目录下特定扩展名的所有文件路径。
    扩展名匹配不区分大小写。
    """
    filepaths = []
    if isinstance(extensions, str):  # 确保 extensions 是元组
        extensions = (extensions,)

    # 将扩展名转换为小写以进行不区分大小写的比较
    lower_extensions = tuple(ext.lower() for ext in extensions)

    for root, _, files in os.walk(directory):
        for file_name in files:  # Renamed 'file' to 'file_name' to avoid conflict with built-in
            if file_name.lower().endswith(lower_extensions):
                filepaths.append(os.path.join(root, file_name))
    general_logger.info(f"在目录 '{directory}' 中找到 {len(filepaths)} 个文件 (扩展名: {extensions})。")
    return filepaths


def get_progress_bar(iterable, desc="Processing", total=None):
    """返回一个配置好的tqdm进度条实例。"""
    if total is None and hasattr(iterable, '__len__'):
        total = len(iterable)
    return tqdm(iterable, desc=desc, total=total, unit="item", ncols=100,
                bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]')


def basic_text_clean(text):
    """基础文本清洗：去除多余空格、换行符，并过滤非标准字符。"""
    if not isinstance(text, str):
        return ""
    # 统一处理换行符为空格
    text = text.replace('\r\n', ' ').replace('\n', ' ').replace('\r', ' ')
    # 替换多个空格为一个空格
    text = re.sub(r'\s+', ' ', text).strip()
    # 仅保留中文、英文、数字、基本标点和一些特定符号
    text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s.,;:!?"\'()《》〈〉“”‘’。，；：！？（）【】「」『』、%+-=@#$^&*]', '', text)
    return text


def split_to_sentences(text, min_len=5):
    """
    将中文文本分割成句子。依赖于常见的中文句末标点。
    """
    if not text:
        return []

    # 规范化一些重复的标点
    text = text.replace("。。", "。").replace("！！", "！").replace("？？", "？")

    # 使用正则表达式按标点分割，同时保留标点符号作为分隔符的一部分
    sentences_parts = re.split(r'([。！？])', text)

    result_sentences = []
    current_sentence_buffer = ""
    for i, part in enumerate(sentences_parts):
        if not part:  # 跳过 re.split 可能产生的空字符串
            continue
        current_sentence_buffer += part
        # 如果part是句末标点，或者这是最后一部分，则处理累积的句子
        if part in '。！？' or i == len(sentences_parts) - 1:
            processed_sentence = current_sentence_buffer.strip()
            if processed_sentence and len(processed_sentence) >= min_len:
                result_sentences.append(processed_sentence)
            current_sentence_buffer = ""  # 重置缓冲区

    # 如果整个文本没有标准句末标点，但长度符合要求，则将其视为一个句子
    if not result_sentences and len(text.strip()) >= min_len:
        result_sentences.append(text.strip())

    return result_sentences


def extract_archive_files(source_dir, target_base_dir, zip_filename_encoding='gbk'):
    """
    解压指定源目录下的所有RAR和ZIP文件到目标目录。
    每个压缩文件解压到目标目录下一个与压缩文件同名（不含扩展名）的子文件夹中。
    对于ZIP文件，会尝试修正中文文件名乱码问题。
    对于RAR文件，文件名编码依赖于unrar工具和系统locale。请确保unrar版本和系统locale配置正确支持中文。

    Args:
        source_dir (str): 包含压缩文件的源目录路径。
        target_base_dir (str): 解压文件的目标基础目录路径。
        zip_filename_encoding (str): ZIP文件中非UTF-8编码的文件名尝试使用的编码，例如 'gbk'。
                                     如果设为 None 或空字符串, 将不进行特殊编码处理。

    Returns:
        bool: 如果所有找到的压缩文件都成功（或尝试）解压，则返回True，否则False。
    """
    if not os.path.isdir(source_dir):
        general_logger.error(f"源目录不存在: {source_dir}")
        return False

    os.makedirs(target_base_dir, exist_ok=True)  # 确保目标基础目录存在

    archive_extensions = ('.rar', '.zip')
    archive_files = get_all_files(source_dir, extensions=archive_extensions)

    if not archive_files:
        general_logger.info(f"在目录 '{source_dir}' 中未找到RAR或ZIP文件。")
        return True  # 没有文件可操作，视为成功完成

    general_logger.info(f"找到 {len(archive_files)} 个压缩文件准备解压。")
    overall_success = True  # 跟踪整体操作是否全部成功

    pbar = get_progress_bar(archive_files, desc="解压压缩文件")
    for archive_filepath in pbar:
        archive_filename = os.path.basename(archive_filepath)
        pbar.set_postfix_str(f"处理: {archive_filename}", refresh=True)

        folder_name = os.path.splitext(archive_filename)[0]
        current_target_dir = os.path.join(target_base_dir, folder_name)
        os.makedirs(current_target_dir, exist_ok=True)

        file_type_processed = ""
        try:
            if archive_filename.lower().endswith('.rar'):
                file_type_processed = "RAR"
                general_logger.info(
                    f"正在解压 {file_type_processed} 文件 '{archive_filepath}' 到 '{current_target_dir}'...")
                with rarfile.RarFile(archive_filepath) as rf:
                    rf.extractall(path=current_target_dir)
                general_logger.info(f"成功解压 {file_type_processed} 文件: {archive_filepath}")

            elif archive_filename.lower().endswith('.zip'):
                file_type_processed = "ZIP"
                general_logger.info(
                    f"正在解压 {file_type_processed} 文件 '{archive_filepath}' 到 '{current_target_dir}'...")
                with zipfile.ZipFile(archive_filepath, 'r') as zf:
                    for member_info in zf.infolist():
                        member_filename_str = member_info.filename

                        if zip_filename_encoding and not (member_info.flag_bits & 0x800):
                            try:
                                original_bytes = member_filename_str.encode(
                                    'cp437')  # Assume zipfile decoded from cp437
                                member_filename_decoded = original_bytes.decode(zip_filename_encoding,
                                                                                'replace')  # 'replace' for robustness
                                general_logger.debug(
                                    f"ZIP成员: '{member_filename_str}' (cp437 bytes: {original_bytes.hex()[:30]}...) "
                                    f"-> 尝试用'{zip_filename_encoding}'解码为: '{member_filename_decoded}'")
                            except Exception as e_decode:  # Catch broader encoding/decoding issues
                                general_logger.warning(
                                    f"无法用'{zip_filename_encoding}'从cp437字节串解码ZIP成员名 '{member_filename_str}'. 错误: {e_decode}. "
                                    f"将使用zipfile库的默认解码结果。")
                                member_filename_decoded = member_filename_str
                        else:
                            member_filename_decoded = member_filename_str
                            if zip_filename_encoding and (member_info.flag_bits & 0x800):
                                general_logger.debug(f"ZIP成员 '{member_filename_decoded}' 已标记为UTF-8。")
                            elif not zip_filename_encoding:
                                general_logger.debug(
                                    f"ZIP成员 '{member_filename_decoded}' 未进行特殊编码处理 (zip_filename_encoding未指定)。")

                        cleaned_path_parts = [part for part in member_filename_decoded.replace('\\', '/').split('/') if
                                              part and part != '.']
                        # Re-check '..' after splitting, as it might be part of a name if not a path segment
                        final_path_parts = []
                        for part in cleaned_path_parts:
                            if part == '..':
                                general_logger.warning(
                                    f"ZIP成员路径中包含 '..' : '{member_filename_decoded}', 将被忽略。")
                            else:
                                final_path_parts.append(part)

                        if not final_path_parts:
                            general_logger.warning(f"ZIP成员路径 '{member_filename_decoded}' 处理后为空，跳过。")
                            continue

                        target_path = os.path.abspath(os.path.join(current_target_dir, *final_path_parts))

                        if not target_path.startswith(os.path.abspath(current_target_dir)):
                            general_logger.error(
                                f"检测到不安全的ZIP成员路径: '{member_filename_decoded}' (解析为 '{target_path}'), "
                                f"试图解压到目标文件夹 '{current_target_dir}' 之外。已跳过此成员。")
                            overall_success = False
                            continue

                        if member_info.is_dir():
                            os.makedirs(target_path, exist_ok=True)
                        else:
                            parent_dir = os.path.dirname(target_path)
                            if parent_dir:
                                os.makedirs(parent_dir, exist_ok=True)
                            try:
                                with zf.open(member_info, 'r') as source_file, \
                                        open(target_path, 'wb') as target_file:
                                    shutil.copyfileobj(source_file, target_file)
                            except Exception as copy_e:
                                general_logger.error(
                                    f"无法写入ZIP成员 '{member_filename_decoded}' 到 '{target_path}': {copy_e}")
                                overall_success = False
                general_logger.info(f"尝试解压 {file_type_processed} 文件 '{archive_filepath}' 中的所有成员。")
            else:
                general_logger.warning(f"跳过未知压缩文件类型: {archive_filepath}")
                continue

        except rarfile.NeedFirstVolume:
            general_logger.error(f"解压RAR失败 (NeedFirstVolume): {archive_filepath}. 这可能是一个分卷压缩文件。")
            overall_success = False
        except (rarfile.NotRarFile, rarfile.BadRarFile) as e_rar:
            general_logger.error(
                f"解压RAR失败 ({type(e_rar).__name__}): {archive_filepath}. 文件可能已损坏或无效. 错误: {e_rar}")
            overall_success = False
        except zipfile.BadZipFile as e_zip:
            general_logger.error(f"解压ZIP失败 (BadZipFile): {archive_filepath}. 文件可能已损坏或无效. 错误: {e_zip}")
            overall_success = False
        except PermissionError as e_perm:
            general_logger.error(
                f"解压失败 (PermissionError): {archive_filepath} -> {current_target_dir}. 错误: {e_perm}")
            overall_success = False
        except Exception as e_unknown:
            general_logger.error(f"解压 '{archive_filepath}' ({file_type_processed}) 时发生未知错误: {e_unknown}")
            overall_success = False
            import traceback
            general_logger.error(traceback.format_exc())

    if overall_success:
        general_logger.info("所有压缩文件解压（或尝试解压所有成员）完成。")
    else:
        general_logger.warning("部分压缩文件或其内部成员解压失败。请检查日志获取详细信息。")

    return overall_success

if __name__ == '__main__':
    # # 测试logger
    test_logger = setup_logger("test_utils_main", "test_utils_main.log", "DEBUG")
    # test_logger.info("这是一条INFO级别的日志。")
    #
    # # --- 文件操作测试 ---
    # temp_dir = os.path.join(PROJECT_ROOT, "temp_test_data_utils")
    # os.makedirs(temp_dir, exist_ok=True)
    # # (之前的测试代码保持不变)
    # temp_text_file = os.path.join(temp_dir, "test.txt")
    # with open(temp_text_file, "w", encoding="utf-8") as f:
    #     f.write("第一行\n")
    #     f.write("第二行 \n")
    #     f.write("\n")
    #     f.write(" 第三行\n")
    #
    # lines = load_text_file(temp_text_file)
    # test_logger.info(f"从 {temp_text_file} 读取的内容: {lines}")
    # assert lines == ["第一行", "第二行", "第三行"]
    #
    # save_text_list(["保存测试1", "保存测试2"], os.path.join(temp_dir, "save_test.txt"))
    # test_json_data = {"name": "测试", "value": 123, "nested": {"key": "嵌套值"}}
    # temp_json_file = os.path.join(temp_dir, "test.json")
    # save_json_file(test_json_data, temp_json_file)
    # loaded_json_data = load_json_file(temp_json_file)
    # test_logger.info(f"从 {temp_json_file} 读取的JSON内容: {loaded_json_data}")
    # assert loaded_json_data == test_json_data
    # sub_temp_dir = os.path.join(temp_dir, "subdir")
    # os.makedirs(sub_temp_dir, exist_ok=True)
    # with open(os.path.join(sub_temp_dir, "sub_test.txt"), "w", encoding="utf-8") as f: f.write("sub dir text")
    # with open(os.path.join(temp_dir, "another.md"), "w", encoding="utf-8") as f: f.write("# Markdown")
    # all_txt_files = get_all_files(temp_dir, extensions=('.txt',))
    # test_logger.info(f"在 {temp_dir} 中找到的txt文件: {all_txt_files}")
    # assert len(all_txt_files) == 2
    # all_files_in_temp = get_all_files(temp_dir, extensions=('.txt', '.md'))
    # test_logger.info(f"在 {temp_dir} 中找到的txt和md文件: {all_files_in_temp}")
    # assert len(all_files_in_temp) == 3
    #
    # # --- 文本清洗和分割测试 ---
    # ocr_sample_text = "  这是\n一个 OCR文本\r\n示例。。它包含一些 混乱格式！  希望能正确 分割？  "
    # cleaned_text = basic_text_clean(ocr_sample_text)
    # test_logger.info(f"原始文本: '{ocr_sample_text}'")
    # test_logger.info(f"清洗后文本: '{cleaned_text}'")
    # sentences = split_to_sentences(cleaned_text, min_len=3)
    # test_logger.info(f"分割后的句子: {sentences}")
    #
    # # --- 压缩文件解压测试 ---
    # test_logger.info("\n======== 测试 RAR 和 ZIP 文件解压 ========")
    # test_archive_source_dir = os.path.join(temp_dir, "archive_source")
    # test_archive_target_dir = os.path.join(temp_dir, "archive_extracted")
    # os.makedirs(test_archive_source_dir, exist_ok=True)
    # os.makedirs(test_archive_target_dir, exist_ok=True)
    #
    # general_logger.info(f"请在以下目录中手动放置一些 .rar 和 .zip 文件以进行解压测试: {test_archive_source_dir}")
    # general_logger.info("例如，创建 test1.rar, test2.zip, subfolder/test3.RAR 等。")
    # general_logger.info("如果该目录下没有压缩文件，解压函数会提示未找到文件。")
    test_archive_source_dir = r'E:\BaiduNetdiskDownload\data\2006-2023CSR文本'
    test_archive_target_dir = r'E:\Datasets\huangmin\CSR'
    # 运行解压函数 (你需要手动放置文件进行测试)
    extract_archive_files(test_archive_source_dir, test_archive_target_dir)
    # general_logger.info(f"测试完成，临时文件在 {temp_dir}。如果需要，请手动删除。")
    try:
        # shutil.rmtree(temp_dir)
        # test_logger.info(f"已删除临时测试目录: {temp_dir}")
        pass # 保持目录以便用户检查
    except OSError as e:
        test_logger.error(f"删除临时目录失败 {temp_dir}: {e}")
    pass
